#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;
class Solution
{
public:
    inline int sign(int num)
    {
        if (num == 0)
        {
            return 0;
        }
        return num > 0 ? 1 : -1;
    }

    /**
     * 时间复杂度：O(n)
     * 间复杂度：O(n)
     */
    int countSubarrays(vector<int> &nums, int k)
    {
        int n = nums.size();
        int kIndex = -1;
        // 先找到index
        for (int i = 0; i < n; i++)
        {
            if (nums[i] == k)
            {
                kIndex = i;
                break;
            }
        }
        int ans = 0;
        unordered_map<int, int> counts;
        counts[0] = 1;
        int sum = 0;
        for (int i = 0; i < n; i++)
        {
            sum += sign(nums[i] - k);
            if (i < kIndex)
            {
                counts[sum]++;
            }
            else
            {
                int prev0 = counts[sum];
                int prev1 = counts[sum - 1];
                ans += prev0 + prev1;
            }
        }
        return ans;
    }

    /**
     * 中位数等于kkk的子数组，那么满足条件的子数组肯定围绕着kkk所在的位置计数。
     * 由于题目知道，大于kkk的与小于kkk的数量差一或者相等，为了方便计数，大于贡献111，小于贡献−1-1−1，本身贡献000.
     * 这样就前缀和的方式计数了，kkk位开始作为结束位，考虑左边起始位的计数，求出和等于当前值与当前值减一的数量，即包含kkk后区间和为0,10,10,1的数量
     */
    int countSubarrays_1(vector<int> &nums, int k)
    {
        unordered_map<int, int> cnt;
        cnt[0] = 1;
        int sum = 0, ans = 0, start = 0;
        for (auto &&x : nums)
        {
            // 大于贡献1，小于贡献−−1，本身贡献0
            sum += (x > k ? 1 : (x == k ? 0 : -1));
            if (start || x == k)
            {
                if (x == k)
                {
                    start = 1;
                }
                ans += cnt[sum] + cnt[sum - 1];
            }
            else
            {
                ++cnt[sum];
            }
        }
        return ans;
    }
};
int main(int argc, char const *argv[])
{
    Solution s;
    vector<int> nums = {3, 4, 1, 2, 5};
    int k = 2;
    cout << s.countSubarrays_1(nums, k) << endl;
    return 0;
}
